# 尝试mujoco显示和控制分开成多进程
# import os
# os.environ["XDG_RUNTIME_DIR"] = f"/run/user/{os.getuid()}"
# os.environ["QT_QPA_PLATFORM"] = "xcb"
import numpy as np
# TODO 尝试一次性发送多条指令
# -*- coding:utf-8 -*-
from ctypes import *
import math
import time

import multiprocessing as mp
from multiprocessing import shared_memory


# 导入自定义函数和CANFD SDK动态链接库>>>
from os.path import dirname, join
import sys
import platform
arch = platform.machine()
sys.path.append(join(dirname(__file__), 'src/lysdemo_'+arch))
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *



def armcontrol(mocap_event, ikq_star_event, share_mocap, share_ikq):

    mocap_pos_quat = np.ndarray(share_mocap_shape, dtype=np.float32, buffer=share_mocap.buf)
    ikq_star = np.ndarray(share_ikq_shape, dtype=np.float32, buffer=share_ikq.buf)


    max_delta_intq_eachloop = 50000
    # >>>>>初始化CAN卡部分>>>
    # 设置can or canfd接口
    lcanlib = LCAN()
    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFDMini, devIndex, canIndex)

    # open device
    ret = lcanlib.LCAN_OpenDevice(LCAN_USBCANFDMini, devIndex)
    if ret != LCAN_STATUS_OK:
        print("LYS USBCANFDMini Open Device failed!")
        exit(0)
    # open success
    print("LYS USBCANFDMini Open Device Success!")

    print('sleeping')
    time.sleep(0.1)

    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFDMini, devIndex, canIndex)

    # 初始化canfd
    ret = lcanlib.LCAN_InitCANFD(LCAN_USBCANFDMini, devIndex, canIndex, canfd_init_cfg)
    if ret != LCAN_STATUS_OK:
        print("LYS USBCANmini Init Failed!")
    else:
        print("LYS USBCANmini Init Success!")
    print('sleeping')
    time.sleep(0.1)

    # 清除can通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFDMini, devIndex, canIndex)
    print('sleeping')
    time.sleep(0.1)
    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    motor_controllers = [MotorController(id) for id in motor_ids]

    for ID in motor_ids:

        # 上使能
        ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFDMini, devIndex, canIndex, motor_controllers[ID-1].can_enable, 1)

        if ret != LCAN_STATUS_OK:
            print("Enable Failed!")
        else:
            print("Enable Success!")
        print('sleeping')
        time.sleep(0.1)


        # 在线更新标志位
        ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFDMini, devIndex, canIndex, motor_controllers[ID-1].can_obj_init, 1)
        
        if ret != LCAN_STATUS_OK:
            print("motor init Failed!")
        else:
            print("motor init Success!")
        print('sleeping')
        time.sleep(0.1)


        # 设置p模式
        ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFDMini, devIndex, canIndex, motor_controllers[ID-1].can_p_mode, 1)
        if ret != LCAN_STATUS_OK:
            print("Set v mode failed!")
        else:
            print("Set v mode success!")

        print('sleeping')
        time.sleep(0.1)

    # 控制
    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFDMini, devIndex, canIndex)

    # 发送与读取>>>
    # P Mode

    pvc = np.zeros((motor_num,3))
    print(f"pvc = {pvc}")
    print(f"pvc[:][0] = {pvc[:][0]}")
    control_time = float(10)
    start = time.time()


    mocap_event.wait() # 等待mujoco给出mocap位置
    mocap_pos = mocap_pos_quat[:3]
    mocap_quat = mocap_pos_quat[3:]

    while time.time() - start < control_time:
        q_real = pvc.T[0]
        q_real_float = q_real*math.pi/180/10000

        # 给定目标点
        mocap_pos = mocap_pos_quat[:3]
        mocap_quat = mocap_pos_quat[3:]
        print(f"get mocap_pos_quat = {mocap_pos_quat}")
        mocap_quat = np.array([1, 0, 0, 0])
        mocap_pos = np.array([.2, 0, 1])

        ikq_star = ik(q_real_float, mocap_pos, mocap_quat).astype(np.float32)
        ikq_star_event.set() # 求出ik_qstar后就通知mujoco

        intikq = trans_q2intq(ikq_star)

        max_delta_intq = np.max(np.abs(intikq - q_real))
        delta_intq_new = (max_delta_intq_eachloop/max(max_delta_intq, max_delta_intq_eachloop))*(intikq - q_real)
        ctrl_intq = np.rint(delta_intq_new + q_real)
        # print(f"ctrl_intq = {ctrl_intq}")
        time.sleep(0.5)
        # last_time = time.time()


        

        for ID in motor_ids:
            # 转换控制数据
            # TODO 单点机控制改多电机一次发送
            p_des = int(ctrl_intq[ID-1])
        
            hex_p = decimal_to_four_hex_numbers(p_des) # TODO 一次生成多个pdes金额hexp

            motor_controllers[ID-1].set_p(hex_p)


            # 发送指令
            # transfer_time1 = time.time()
            ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFDMini, devIndex, canIndex, motor_controllers[ID-1].can_set_p, 1)
            # print(f"transfer time cost: {time.time() - transfer_time1}")
            # if ret != LCAN_STATUS_OK:
            #     print(f"Set {ID} p failed!")
            # else:
            #     print(f"Set {ID} p success!")


        # 读取状态
        new_pvc = receive_pvc(lcanlib, devIndex, canIndex)
        # loop_time = time.time() - last_time
        if new_pvc is not None:
            pvc = new_pvc  # 仅当有有效数据时更新
            # print(f"最新状态: {pvc}")
            # print(f"time cost: {loop_time}")
        else:
            print(f"get no pvc data\n")


        
    # 发送与读取<<<

    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFDMini, devIndex, canIndex)
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    #Close Device
    lcanlib.LCAN_CloseDevice(LCAN_USBCANFDMini, devIndex)
    print("LYS USBCANmini Close!")

    # 清理共享内存
    share_mocap.close()
    share_mocap.unlink()
    share_ikq.close()
    share_ikq.unlink()


def process_test(mocap_event, ikq_star_event, share_mocap, share_ikq):
    
    mocap_pos_quat = np.ndarray(share_mocap_shape, dtype=np.float32, buffer=share_mocap.buf)
    ikq_star = np.ndarray(share_ikq_shape, dtype=np.float32, buffer=share_ikq.buf)
    while True:
        ikq_star = np.array([.1,.1,.1,.1,.1,.1]).astype(np.float32)
        print(f"ikq_star test = {ikq_star}")
        ikq_star_event.set()
        time.sleep(0.1)
        


def mujoco_show(mocap_event, ikq_star_event, share_mocap, share_ikq):
    import mujoco
    import mujoco.viewer as viewer
    print("Start mujoco_show")
    mocap_pos_quat = np.ndarray(share_mocap_shape, dtype=np.float32, buffer=share_mocap.buf)
    ikq_star = np.ndarray(share_ikq_shape, dtype=np.float32, buffer=share_ikq.buf)

    # mujoco加载模型
    model = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene_P.xml')) 
    data = mujoco.MjData(model)
    mocap_id = model.body("target").mocapid[0]
    print("model load success")
    with viewer.launch_passive(model,data) as viewer:
        print("viewer launch success")
        # Initialize the camera view to that of the free camera.
        mujoco.mjv_defaultFreeCamera(model, viewer.cam)
        print("camera init success")
        viewer.opt.frame = mujoco.mjtFrame.mjFRAME_SITE
        print("frame init success")
        while viewer.is_running():


            mocap_pos = data.mocap_pos[mocap_id]
            mocap_quat = data.mocap_quat[mocap_id]
            mocap_pos_quat = np.concatenate((mocap_pos, mocap_quat)).astype(np.float32)
            mocap_event.set() # 得到mocap位置后通知control
            time.sleep(0.1)

            ikq_star_event.wait() # 等待返回的ik_qstar

            print("ikq_star_event success")
            data.qpos = ikq_star
            # print(f"data.qpos: {data.qpos}")
            print(f"ikq_star = {ikq_star}")
            
            mujoco.mj_step(model, data)
            mujoco.mj_forward(model, data)
            

            viewer.sync()
            time.sleep(model.opt.timestep)

    # 清理共享内存
    share_mocap.close()
    share_mocap.unlink()
    share_ikq.close()
    share_ikq.unlink()


if __name__ == "__main__":
    share_mocap_shape = (7,)
    share_ikq_shape = (6,)
    share_mocap = shared_memory.SharedMemory(create=True, size=np.prod(share_mocap_shape)*4)  # float32
    share_ikq = shared_memory.SharedMemory(create=True, size=np.prod(share_ikq_shape)*4)  # float32
    mocap_event = mp.Event()
    ikq_star_event = mp.Event()

    

    # p_test = mp.Process(target=process_test, args=(mocap_event, ikq_star_event, share_mocap, share_ikq))
    # p_test.start()
    # p_control = mp.Process(target=armcontrol, args=(mocap_event, ikq_star_event))
    # p_control.start()
    # mujoco_show(mocap_event, ikq_star_event, share_mocap, share_ikq)

    # p_control.join()
    # p_mujoco.terminate()
